package com.thiswhyme.elasticsearch.core;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.thiswhyme.common.expection.BusinessException;
import com.thiswhyme.common.response.ResponseCode;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

/**
 * @description: 基础es文档操作客户端
 * @author: thiswhyme
 * @date: 2021/10/14
 */
@Slf4j
@Component
@AllArgsConstructor
public class BaseDocumentClient {

    private final RestHighLevelClient client;

    /**
     * 新增文档
     * @param document
     * @param <T>
     */
    public <T extends BaseDocument> void insert(T document) {
        IndexRequest request = new IndexRequest(document.getIndex())
                .id(document.getId())
                .create(true)
                .source(document.toJson(), XContentType.JSON);
        try {
            IndexResponse response = client.index(request, RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档新增完成，结果：{}", response.getResult());
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档新增失败，index：{}， id：{}", document.getIndex(), document.getId(), e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档新增失败");
        }
    }

    /**
     * 批量新增文档
     * @param documents
     * @param <T>
     */
    public <T extends BaseDocument> void batchInsert(List<T> documents) {
        BulkRequest bulkRequest = new BulkRequest();
        documents.forEach(document -> {
            IndexRequest request = new IndexRequest(document.getIndex())
                    .id(document.getId())
                    .create(true)
                    .source(document.toJson(), XContentType.JSON);
            bulkRequest.add(request);
        });
        try {
            BulkResponse response = client.bulk(bulkRequest, RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档批量新增完成，结果：{}", getBulkResults(response));
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档批量新增失败", e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档批量新增失败");
        }
    }

    /**
     * 更新文档
     * @param <T>
     * @param index
     * @param update
     */
    public <T extends BaseDocument> void update(String index, DocUpdateVO update) {
        UpdateRequest updateRequest = new UpdateRequest(index, update.getId()).doc(update.getUpdates());
        try {
            UpdateResponse response = client.update(updateRequest, RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档更新完成，结果：{}", response.getResult());
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档更新失败，index：{}， id：{}", index, update.getId(), e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档更新失败");
        }
    }

    /**
     * 批量更新文档
     * @param <T>
     * @param index
     * @param updates
     */
    public <T extends BaseDocument> void batchUpdate(String index, List<DocUpdateVO> updates) {
        BulkRequest bulkRequest = new BulkRequest();
        updates.forEach(update -> {
            UpdateRequest request = new UpdateRequest(index, update.getId())
                    .doc(update.getUpdates());
            bulkRequest.add(request);
        });
        try {
            BulkResponse response = client.bulk(bulkRequest, RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档批量更新完成，结果：{}", getBulkResults(response));
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档批量更新失败，index：{}", index, e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档批量更新失败");
        }
    }

    /**
     * 删除文档
     * @param index
     * @param id
     * @return
     */
    public void delete(String index, String id) {
        try {
            DeleteResponse response = client.delete(new DeleteRequest(index, id), RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档删除完成，结果：{}", response.getResult());
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档删除异常，index：{}， id：{}", index, id, e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档删除异常");
        }
    }

    /**
     * 批量删除文档
     * @param index
     * @param ids
     */
    public void batchDelete(String index, List<String> ids) {
        BulkRequest bulkRequest = new BulkRequest();
        ids.forEach(id -> {
            DeleteRequest request = new DeleteRequest(index, id);
            bulkRequest.add(request);
        });
        try {
            BulkResponse response = client.bulk(bulkRequest, RequestOptions.DEFAULT);
            if (log.isDebugEnabled()) {
                log.debug("文档批量删除完成，结果：{}", getBulkResults(response));
            }
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档批量删除失败，index：{}", index, e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档批量删除失败");
        }
    }

    /**
     * 获取批处理结果
     * @param response
     * @return
     */
    private List<String> getBulkResults(BulkResponse response) {
        List<String> results = new ArrayList<>();
        for (BulkItemResponse item : response.getItems()) {
            if (null != item.getResponse()) {
                results.add(item.getResponse().getResult().name());
            }
            if (StringUtils.hasText(item.getFailureMessage())) {
                results.add(item.getFailureMessage());
            }
        }
        return results;
    }

    /**
     * 搜索
     *
     * @param queryBuilder
     * @param clazz
     * @param indices
     * @param <T>
     * @return
     */
    public <T extends BaseDocument> List<T> search(QueryBuilder queryBuilder, Class<T> clazz, String[] indices) {
        return this.search(queryBuilder, null, clazz, indices);
    }

    /**
     * 搜索-排序
     *
     * @param queryBuilder
     * @param clazz
     * @param indices
     * @param <T>
     * @return
     */
    public <T extends BaseDocument> List<T> search(QueryBuilder queryBuilder, SortBuilder sortBuilder, Class<T> clazz, String[] indices) {
        return this.search(queryBuilder, sortBuilder, null, null, clazz, indices);
    }

    /**
     * 搜索-分页&排序
     *
     * @param queryBuilder
     * @param sortBuilder
     * @param page
     * @param size
     * @param clazz
     * @param indices
     * @return
     * @param <T>
     */
    public <T extends BaseDocument> List<T> search(QueryBuilder queryBuilder, SortBuilder sortBuilder, Integer page,
                                                   Integer size, Class<T> clazz, String[] indices) {

        // 构造SearchSourceBuilder
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(queryBuilder);
        Optional.ofNullable(sortBuilder).ifPresent(builder -> searchSourceBuilder.sort(sortBuilder));
        if (null != page) {
            Assert.isTrue(page >= 1, "页码必须大于等于1");
            size = Optional.ofNullable(size).orElse(20);
            searchSourceBuilder.from((page - 1) * size);
            searchSourceBuilder.size(size);
        }

        // 构造SearchRequest
        SearchRequest searchRequest = new SearchRequest(indices);
        searchRequest.source(searchSourceBuilder);

        // 执行请求
        SearchResponse searchResponse;
        try {
            searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档查询异常，queryBuilder：{}， indices：{}", JSONUtil.toJsonStr(queryBuilder), indices, e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档分页查询异常");
        }
        SearchHits hits = searchResponse.getHits();
        SearchHit[] searchHits = hits.getHits();

        List<T> list = CollUtil.newArrayList();
        for (SearchHit hit : searchHits) {
            T t = JSONUtil.toBean(hit.getSourceAsString(), clazz);
            // 数据的唯一标识
            t.setId(hit.getId());
            // 索引
            t.setIndex(hit.getIndex());
            // 排序值
            t.setSortValue(hit.getSortValues());
            list.add(t);
        }
        return list;
    }

    /**
     * 查询数量
     *
     * @param queryBuilder
     * @param indices
     * @return
     */
    public long count(QueryBuilder queryBuilder, String[] indices) {
        // 构造请求
        CountRequest countRequest = new CountRequest(indices);
        countRequest.query(queryBuilder);

        // 执行请求
        CountResponse countResponse;
        try {
            countResponse = client.count(countRequest, RequestOptions.DEFAULT);
        } catch (IOException | ElasticsearchStatusException e) {
            log.error("文档数量查询异常，queryBuilder：{}， indices：{}", JSONUtil.toJsonStr(queryBuilder), indices, e);
            throw new BusinessException(ResponseCode.OPERATION_FAILURE.getCode(), "文档数量查询异常");
        }
        return countResponse.getCount();
    }
}
